{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "dc601a8c-75a7-41c0-8fc2-e0a4edbc6284",
   "metadata": {},
   "outputs": [
    {
     "ename": "ModuleNotFoundError",
     "evalue": "No module named 'tushare'",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mModuleNotFoundError\u001b[0m                       Traceback (most recent call last)",
      "Cell \u001b[0;32mIn[2], line 1\u001b[0m\n\u001b[0;32m----> 1\u001b[0m \u001b[38;5;28;01mimport\u001b[39;00m \u001b[38;5;21;01mtushare\u001b[39;00m \u001b[38;5;28;01mas\u001b[39;00m \u001b[38;5;21;01mts\u001b[39;00m\n\u001b[1;32m      2\u001b[0m \u001b[38;5;28;01mimport\u001b[39;00m \u001b[38;5;21;01mnumpy\u001b[39;00m \u001b[38;5;28;01mas\u001b[39;00m \u001b[38;5;21;01mnp\u001b[39;00m\n\u001b[1;32m      3\u001b[0m \u001b[38;5;28;01mimport\u001b[39;00m \u001b[38;5;21;01mpandas\u001b[39;00m \u001b[38;5;28;01mas\u001b[39;00m \u001b[38;5;21;01mpd\u001b[39;00m\n",
      "\u001b[0;31mModuleNotFoundError\u001b[0m: No module named 'tushare'"
     ]
    }
   ],
   "source": [
    "import tushare as ts\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import matplotlib.pyplot as plt\n",
    "import seaborn as sns\n",
    "from datetime import datetime, timedelta\n",
    "from sklearn.preprocessing import StandardScaler, MinMaxScaler\n",
    "from sklearn.decomposition import PCA\n",
    "from sklearn.impute import SimpleImputer\n",
    "from sklearn.model_selection import train_test_split\n",
    "from sklearn.ensemble import RandomForestClassifier\n",
    "from sklearn.metrics import classification_report, confusion_matrix, roc_curve, auc\n",
    "from sklearn.utils import resample\n",
    "\n",
    "# 设置Tushare token\n",
    "ts.set_token('1c7f85b9026518588c0d0cdac712c2d17344332c9c8cfe6bc83ee75c')\n",
    "pro = ts.pro_api()\n",
    "\n",
    "\n",
    "# 1. 获取两年至少10支股票的日线行情数据\n",
    "def get_stock_data():\n",
    "    # 获取当前日期和两年前的日期\n",
    "    end_date = datetime.now().strftime('%Y%m%d')\n",
    "    start_date = (datetime.now() - timedelta(days=730)).strftime('%Y%m%d')\n",
    "\n",
    "    # 选择10只股票代码 (这里选择了一些大盘股作为示例)\n",
    "    stock_codes = [\n",
    "        '600519.SH',  # 贵州茅台\n",
    "        '601318.SH',  # 中国平安\n",
    "        '000858.SZ',  # 五粮液\n",
    "        '600036.SH',  # 招商银行\n",
    "        '600030.SH',  # 中信证券\n",
    "        '000001.SZ',  # 平安银行\n",
    "        '601398.SH',  # 工商银行\n",
    "        '601988.SH',  # 中国银行\n",
    "        '600028.SH',  # 中国石化\n",
    "        '601857.SH'  # 中国石油\n",
    "    ]\n",
    "\n",
    "    # 存储所有股票数据\n",
    "    all_data = pd.DataFrame()\n",
    "\n",
    "    for code in stock_codes:\n",
    "        try:\n",
    "            # 获取日线数据\n",
    "            df = pro.daily(ts_code=code, start_date=start_date, end_date=end_date)\n",
    "            df = df.sort_values('trade_date')\n",
    "            df['ts_code'] = code\n",
    "            all_data = pd.concat([all_data, df])\n",
    "        except Exception as e:\n",
    "            print(f\"Error fetching data for {code}: {e}\")\n",
    "\n",
    "    # 转换日期格式并设置为索引\n",
    "    all_data['trade_date'] = pd.to_datetime(all_data['trade_date'])\n",
    "    all_data = all_data.set_index('trade_date').sort_index()\n",
    "\n",
    "    return all_data\n",
    "\n",
    "# 获取数据\n",
    "stock_data = get_stock_data()\n",
    "print(f\"获取到的数据形状: {stock_data.shape}\")\n",
    "print(stock_data.head())\n",
    "\n",
    "# 2. 数据分类(打标签)与技术指标计算\n",
    "def preprocess_data(df):\n",
    "    # 创建收益率标签 (明天上涨为1，下跌为0)\n",
    "    df['next_close'] = df.groupby('ts_code')['close'].shift(-1)\n",
    "    df['label'] = np.where(df['next_close'] > df['close'], 1, 0)\n",
    "    # 计算技术指标 (超过15个)\n",
    "    # 1. 简单移动平均线\n",
    "    df['MA5'] = df.groupby('ts_code')['close'].transform(lambda x: x.rolling(5).mean())\n",
    "    df['MA10'] = df.groupby('ts_code')['close'].transform(lambda x: x.rolling(10).mean())\n",
    "    df['MA20'] = df.groupby('ts_code')['close'].transform(lambda x: x.rolling(20).mean())\n",
    "    # 2. 指数移动平均线\n",
    "    df['EMA5'] = df.groupby('ts_code')['close'].transform(lambda x: x.ewm(span=5).mean())\n",
    "    df['EMA10'] = df.groupby('ts_code')['close'].transform(lambda x: x.ewm(span=10).mean())\n",
    "    # 3. MACD\n",
    "    df['EMA12'] = df.groupby('ts_code')['close'].transform(lambda x: x.ewm(span=12).mean())\n",
    "    df['EMA26'] = df.groupby('ts_code')['close'].transform(lambda x: x.ewm(span=26).mean())\n",
    "    df['MACD'] = df['EMA12'] - df['EMA26']\n",
    "    df['Signal_Line'] = df['MACD'].ewm(span=9).mean()\n",
    "    df['MACD_Hist'] = df['MACD'] - df['Signal_Line']\n",
    "    # 4. RSI\n",
    "    delta = df.groupby('ts_code')['close'].diff(1)\n",
    "    gain = delta.where(delta > 0, 0)\n",
    "    loss = -delta.where(delta < 0, 0)\n",
    "    avg_gain = gain.rolling(14).mean()\n",
    "    avg_loss = loss.rolling(14).mean()\n",
    "    df['RS'] = avg_gain / avg_loss\n",
    "    df['RSI'] = 100 - (100 / (1 + df['RS']))\n",
    "    # 5. 布林带\n",
    "    df['Middle_Band'] = df.groupby('ts_code')['close'].transform(lambda x: x.rolling(20).mean())\n",
    "    df['Upper_Band'] = df['Middle_Band'] + 2 * df.groupby('ts_code')['close'].transform(\n",
    "        lambda x: x.rolling(20).std())\n",
    "    df['Lower_Band'] = df['Middle_Band'] - 2 * df.groupby('ts_code')['close'].transform(\n",
    "        lambda x: x.rolling(20).std())\n",
    "    # 6. KDJ\n",
    "    low_min = df.groupby('ts_code')['low'].transform(lambda x: x.rolling(9).min())\n",
    "    high_max = df.groupby('ts_code')['high'].transform(lambda x: x.rolling(9).max())\n",
    "    rsv = (df['close'] - low_min) / (high_max - low_min) * 100\n",
    "    df['K'] = rsv.ewm(com=2).mean()\n",
    "    df['D'] = df['K'].ewm(com=2).mean()\n",
    "    df['J'] = 3 * df['K'] - 2 * df['D']\n",
    "    # 7. 成交量相关指标\n",
    "    df['Vol_MA5'] = df.groupby('ts_code')['vol'].transform(lambda x: x.rolling(5).mean())\n",
    "    df['Vol_MA10'] = df.groupby('ts_code')['vol'].transform(lambda x: x.rolling(10).mean())\n",
    "    df['Vol_RSI'] = 100 * df['vol'] / (df['Vol_MA5'] + 1e-6)  # 防止除以0\n",
    "    # 8. 价格波动率\n",
    "    df['Price_Vol'] = df.groupby('ts_code')['close'].transform(lambda x: x.rolling(5).std() / x.rolling(5).mean())\n",
    "    # 9. 动量指标\n",
    "    df['Momentum_5'] = df['close'] / df.groupby('ts_code')['close'].shift(5) - 1\n",
    "    df['Momentum_10'] = df['close'] / df.groupby('ts_code')['close'].shift(10) - 1\n",
    "    # 10. 收益率\n",
    "    df['Return_1'] = df['close'].pct_change()\n",
    "    df['Return_5'] = df['close'] / df.groupby('ts_code')['close'].shift(5) - 1\n",
    "    df['Return_10'] = df['close'] / df.groupby('ts_code')['close'].shift(10) - 1\n",
    "    # 删除包含NaN的行\n",
    "    df = df.dropna()\n",
    "    return df\n",
    "\n",
    "processed_data = preprocess_data(stock_data)\n",
    "print(f\"处理后的数据形状: {processed_data.shape}\")\n",
    "print(processed_data.head())\n",
    "\n",
    "# 3. 建模前的数据处理与分析\n",
    "def prepare_data_for_modeling(df):\n",
    "    # 选择特征和标签\n",
    "    features = [\n",
    "        'open', 'high', 'low', 'close', 'vol', 'MA5', 'MA10', 'MA20',\n",
    "        'EMA5', 'EMA10', 'MACD', 'Signal_Line', 'MACD_Hist', 'RSI',\n",
    "        'Middle_Band', 'Upper_Band', 'Lower_Band', 'K', 'D', 'J',\n",
    "        'Vol_MA5', 'Vol_MA10', 'Vol_RSI', 'Price_Vol', 'Momentum_5',\n",
    "        'Momentum_10', 'Return_1', 'Return_5', 'Return_10'\n",
    "    ]\n",
    "\n",
    "    # 添加股票代码作为分类特征\n",
    "    features.append('ts_code')\n",
    "\n",
    "    X = df[features]\n",
    "    y = df['label']\n",
    "\n",
    "    # 处理分类特征 (股票代码)\n",
    "    X = pd.get_dummies(X, columns=['ts_code'])\n",
    "\n",
    "    # 空值处理\n",
    "    imputer = SimpleImputer(strategy='median')\n",
    "    X = pd.DataFrame(imputer.fit_transform(X), columns=X.columns)\n",
    "\n",
    "    # 异常值处理 - 使用IQR方法\n",
    "    for col in X.select_dtypes(include=[np.number]).columns:\n",
    "        Q1 = X[col].quantile(0.25)\n",
    "        Q3 = X[col].quantile(0.75)\n",
    "        IQR = Q3 - Q1\n",
    "        lower_bound = Q1 - 1.5 * IQR\n",
    "        upper_bound = Q3 + 1.5 * IQR\n",
    "        X[col] = np.where(X[col] < lower_bound, lower_bound, X[col])\n",
    "        X[col] = np.where(X[col] > upper_bound, upper_bound, X[col])\n",
    "\n",
    "    # 归一化\n",
    "    scaler = MinMaxScaler()\n",
    "    X_scaled = pd.DataFrame(scaler.fit_transform(X), columns=X.columns)\n",
    "\n",
    "    # 主成分分析 (PCA)\n",
    "    pca = PCA(n_components=0.95)  # 保留95%的方差\n",
    "    X_pca = pca.fit_transform(X_scaled)\n",
    "    print(f\"PCA降维后的维度: {X_pca.shape[1]}\")\n",
    "\n",
    "    # 可视化PCA结果\n",
    "    plt.figure(figsize=(10, 6))\n",
    "    plt.bar(range(1, len(pca.explained_variance_ratio_) + 1), pca.explained_variance_ratio_,\n",
    "            alpha=0.5, align='center', label='Individual explained variance')\n",
    "    plt.step(range(1, len(pca.explained_variance_ratio_) + 1), np.cumsum(pca.explained_variance_ratio_),\n",
    "             where='mid', label='Cumulative explained variance')\n",
    "    plt.ylabel('Explained variance ratio')\n",
    "    plt.xlabel('Principal components')\n",
    "    plt.legend(loc='best')\n",
    "    plt.title('PCA Explained Variance')\n",
    "    plt.show()\n",
    "\n",
    "    # 相关性分析\n",
    "    corr_matrix = X_scaled.corr()\n",
    "    plt.figure(figsize=(20, 16))\n",
    "    sns.heatmap(corr_matrix, annot=False, cmap='coolwarm', center=0)\n",
    "    plt.title('Feature Correlation Matrix')\n",
    "    plt.show()\n",
    "\n",
    "    # 个股画像 (选择6个属性)\n",
    "    # 这里我们选择一些关键的技术指标作为个股画像\n",
    "    profile_features = ['close', 'vol', 'RSI', 'MACD_Hist', 'Momentum_5', 'Return_1']\n",
    "\n",
    "    # 数据均衡处理\n",
    "    # 先分割数据\n",
    "    X_train, X_test, y_train, y_test = train_test_split(X_scaled, y, test_size=0.3, random_state=42)\n",
    "    X_train = X_train.reset_index(drop=True)\n",
    "    y_train = y_train.reset_index(drop=True)\n",
    "\n",
    "    # 检查类别分布\n",
    "    print(\"训练集类别分布:\")\n",
    "    print(y_train.value_counts())\n",
    "    print(\"\\n测试集类别分布:\")\n",
    "    print(y_test.value_counts())\n",
    "\n",
    "    # 可视化类别分布\n",
    "    plt.figure(figsize=(8, 5))\n",
    "    sns.countplot(x=y_train)\n",
    "    plt.title('Class Distribution in Training Set')\n",
    "    plt.show()\n",
    "\n",
    "    # 使用随机过采样\n",
    "    # 分离多数类和少数类\n",
    "    df_majority = X_train[y_train == 0]\n",
    "    df_minority = X_train[y_train == 1]\n",
    "\n",
    "    # 对少数类进行过采样\n",
    "    df_minority_upsampled = resample(df_minority,\n",
    "                                     replace=True,  # 有放回抽样\n",
    "                                     n_samples=len(df_majority),  # 调整到与多数类相同的样本数\n",
    "                                     random_state=42)\n",
    "\n",
    "    # 合并多数类和过采样后的少数类\n",
    "    X_train_res = pd.concat([df_majority, df_minority_upsampled])\n",
    "    y_train_res = pd.concat([y_train[y_train == 0], pd.Series([1] * len(df_minority_upsampled))])\n",
    "\n",
    "    # 打乱数据\n",
    "    indices = np.arange(len(X_train_res))\n",
    "    np.random.shuffle(indices)\n",
    "    X_train_res = X_train_res.iloc[indices]\n",
    "    y_train_res = y_train_res.iloc[indices]\n",
    "\n",
    "    print(\"\\n过采样后训练集类别分布:\")\n",
    "    print(y_train_res.value_counts())\n",
    "\n",
    "    return X_train_res, X_test, y_train_res, y_test, profile_features\n",
    "\n",
    "X_train_res, X_test, y_train_res, y_test, profile_features = prepare_data_for_modeling(processed_data)\n",
    "\n",
    "\n",
    "# 4. 机器学习建模与评估\n",
    "def build_and_evaluate_model(X_train, X_test, y_train, y_test):\n",
    "    # 初始化随机森林分类器\n",
    "    model = RandomForestClassifier(n_estimators=100, random_state=42)\n",
    "\n",
    "    # 训练模型\n",
    "    model.fit(X_train, y_train)\n",
    "\n",
    "    # 预测\n",
    "    y_pred = model.predict(X_test)\n",
    "    y_pred_proba = model.predict_proba(X_test)[:, 1]\n",
    "\n",
    "    # 评估模型\n",
    "    print(\"\\n模型评估报告:\")\n",
    "    print(classification_report(y_test, y_pred))\n",
    "\n",
    "    # 混淆矩阵\n",
    "    cm = confusion_matrix(y_test, y_pred)\n",
    "    plt.figure(figsize=(6, 4))\n",
    "    sns.heatmap(cm, annot=True, fmt='d', cmap='Blues',\n",
    "                xticklabels=['Predicted 0', 'Predicted 1'],\n",
    "                yticklabels=['Actual 0', 'Actual 1'])\n",
    "    plt.title('Confusion Matrix')\n",
    "    plt.show()\n",
    "\n",
    "    # ROC曲线\n",
    "    fpr, tpr, thresholds = roc_curve(y_test, y_pred_proba)\n",
    "    roc_auc = auc(fpr, tpr)\n",
    "\n",
    "    plt.figure(figsize=(8, 6))\n",
    "    plt.plot(fpr, tpr, color='darkorange', lw=2,\n",
    "             label=f'ROC curve (area = {roc_auc:.2f})')\n",
    "    plt.plot([0, 1], [0, 1], color='navy', lw=2, linestyle='--')\n",
    "    plt.xlim([0.0, 1.0])\n",
    "    plt.ylim([0.0, 1.05])\n",
    "    plt.xlabel('False Positive Rate')\n",
    "    plt.ylabel('True Positive Rate')\n",
    "    plt.title('Receiver Operating Characteristic')\n",
    "    plt.legend(loc=\"lower right\")\n",
    "    plt.show()\n",
    "    # 特征重要性\n",
    "    feature_importance = model.feature_importances_\n",
    "    feature_names = X_train.columns\n",
    "    # 选择最重要的20个特征\n",
    "    important_indices = np.argsort(feature_importance)[-20:]\n",
    "    important_features = [feature_names[i] for i in important_indices]\n",
    "    important_importances = feature_importance[important_indices]\n",
    "    plt.figure(figsize=(12, 8))\n",
    "    sns.barplot(x=important_importances, y=important_features, palette='viridis')\n",
    "    plt.title('Top 20 Feature Importances')\n",
    "    plt.xlabel('Importance')\n",
    "    plt.ylabel('Features')\n",
    "    plt.show()\n",
    "    return model\n",
    "\n",
    "# 构建并评估模型\n",
    "model = build_and_evaluate_model(X_train_res, X_test, y_train_res, y_test)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python [conda env:base] *",
   "language": "python",
   "name": "conda-base-py"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
